import os

# 保证种子文件夹存在
SEED_DIR = "."

OUTPUT_FILE = "my_model_seeds.py"

# 模板代码：为每个 seed 生成一个 get_model_X 函数
TEMPLATE = """
def get_model_{i}():
    import tvm
    with open("{seed_path}", "r") as f:
        return tvm.script.from_source(f.read()), None
"""

def main():
    # 获取所有以数字命名的 txt 文件
    seed_files = sorted(f for f in os.listdir(SEED_DIR) if f.endswith(".txt") and f[:-4].isdigit())
    with open(OUTPUT_FILE, "w") as out:
        out.write("# Auto-generated model seed loader functions\n")
        out.write("# Each function returns (relax.IRModule, None)\n\n")
        for fname in seed_files:
            i = fname[:-4]  # 去掉 ".txt"
            seed_path = os.path.join(SEED_DIR, fname)
            out.write(TEMPLATE.format(i=i, seed_path=seed_path.replace("\\", "/")))

    print(f"✅ 生成完成，共生成 {len(seed_files)} 个函数，保存至 {OUTPUT_FILE}")

if __name__ == "__main__":
    main()
